package gsthrottle

import (
	"context"
	"sync"
	"time"
)

type TokenThrottle struct {
	bucket       map[int64]uint
	maxBucket    map[int64]uint
	rate         map[int64]uint
	rateTimeUnit time.Duration
	ctx          context.Context
	cancel       context.CancelFunc
	sync.Mutex
	user int64
}

func (x *TokenThrottle) AddThrottle(throttleConfig ...*ThrottleConfig) {
	x.Lock()
	defer x.Unlock()

	for _, v := range throttleConfig {
		if _, ok := x.bucket[v.Id]; ok {
			continue
		}

		x.bucket[v.Id] = 0
		x.rate[v.Id] = v.Limit
		x.maxBucket[v.Id] = v.MaxBucket
	}
}

func (x *TokenThrottle) UpdateThrottle(throttleConfig *ThrottleConfig) {
	x.Lock()
	defer x.Unlock()

	if _, ok := x.bucket[throttleConfig.Id]; ok {
		x.rate[throttleConfig.Id] = throttleConfig.Limit
		x.maxBucket[throttleConfig.Id] = throttleConfig.MaxBucket
	}
}

func (x *TokenThrottle) RemoveId(id int64) bool {
	x.Lock()
	defer x.Unlock()

	if _, ok := x.bucket[id]; ok {
		delete(x.bucket, id)
		return true
	}

	return false

}

func (x *TokenThrottle) Init() {
	x.ctx, x.cancel = context.WithCancel(context.Background())
	x.pushBucket()
	go func() {
		// 按速率生成bucket
		ticker := time.NewTicker(x.rateTimeUnit)
		defer ticker.Stop()
		for {
			select {
			case <-x.ctx.Done():
				return
			case <-ticker.C:
				x.pushBucket()
			}
		}
	}()
}

func (x *TokenThrottle) pushBucket() {
	x.Lock()
	defer x.Unlock()

	for k, _ := range x.bucket {
		if x.bucket[k] < x.maxBucket[k] {
			x.bucket[k] += x.rate[k]
		}
	}
}

func (x *TokenThrottle) Call(id int64, call func(), reject func()) {
	x.Lock()
	defer x.Unlock()

	if x.bucket[id] <= 0 {
		reject()
		return
	}

	x.bucket[id]--
	call()
}

func (x *TokenThrottle) Close() {
	x.cancel()
}
